import pandas as pd
import numpy as np
import json
import random
import os
from data_aug import resize
import cv2



def resize_anno(train_anno_file_dir,val_anno_file_dir,train_anno_save_dir,val_anno_save_dir,image_save_dir,size=(1152,1536)):
    def fun(anno_file_dir,anno_save_dir,image_save_dir,size):
        if not os.path.exists(image_save_dir):
            os.makedirs(image_save_dir)
        anno_file = pd.read_csv(anno_file_dir)
        for i in range(len(anno_file)):
            image_dir = anno_file.iloc[i,0]
            image = cv2.imread(image_dir)
            bboxex = [[anno_file.iloc[i,1],anno_file.iloc[i,2],anno_file.iloc[i,3],anno_file.iloc[i,4]]]
            category_id = anno_file.iloc[i,5]
            print(bboxex,category_id)
            image,bboxex,_ = resize(image=image,bboxes=bboxex,category_id=category_id,size=size)
            resized_image_dir = image_dir.replace('train',image_save_dir.split('data/')[1])
            if not os.path.exists(resized_image_dir):
                cv2.imwrite(resized_image_dir,image)
            anno_file.iloc[i,0] = resized_image_dir
            anno_file.iloc[i,1],anno_file.iloc[i,2],anno_file.iloc[i,3],anno_file.iloc[i,4], = int(bboxex[0][0]),int(bboxex[0][1]),int(bboxex[0][2]),int(bboxex[0][3])
        anno_file.to_csv(anno_save_dir,index=False)
    fun(train_anno_file_dir,train_anno_save_dir,image_save_dir,size=size)
    fun(val_anno_file_dir,val_anno_save_dir,image_save_dir,size=size)

if __name__ == '__main__':

    resize_anno(train_anno_file_dir='../data/annotations_train_b.csv',
                val_anno_file_dir='../data/annotations_val_b.csv',
                train_anno_save_dir='../data/annotations_train_b_1152_1536.csv',
                val_anno_save_dir='../data/annotations_val_b_1152_1536.csv',
                image_save_dir='../data/resized_train_b/1152_1536')





